#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>

#include <cassert>
#include <vector>

#include "compat.h"

void compute_n1_n2(at::Tensor input, at::IntArrayRef normalized_shape, int& n1, int& n2) {
    int idiff = input.ndimension() - normalized_shape.size();
    n2 = 1;
    for (int i = 0; i < (int)normalized_shape.size(); ++i) {
        assert(input.sizes()[i + idiff] == normalized_shape[i]);
        n2 *= normalized_shape[i];
    }
    n1 = 1;
    for (int i = 0; i < idiff; ++i) {
        n1 *= input.sizes()[i];
    }
}

void check_args(at::IntArrayRef normalized_shape, at::Tensor gamma, at::Tensor beta) {
    TORCH_CHECK(!gamma.defined() || gamma.sizes().equals(normalized_shape));
    TORCH_CHECK(!beta.defined() || beta.sizes().equals(normalized_shape));
}

void check_args(at::Tensor input, at::IntArrayRef normalized_shape, int& n1, int& n2) {
    int64_t normalized_ndim = normalized_shape.size();

    if (normalized_ndim < 1) {
        std::stringstream ss;
        ss << "Expected normalized_shape to be at least 1-dimensional, i.e., "
           << "containing at least one element, but got normalized_shape=" << normalized_shape;
        throw std::runtime_error(ss.str());
    }

    auto input_shape = input.sizes();
    auto input_ndim = input.dim();

    if (input_ndim < normalized_ndim ||
        !input_shape.slice(input_ndim - normalized_ndim).equals(normalized_shape)) {
        std::stringstream ss;
        ss << "Given normalized_shape=" << normalized_shape << ", expected input with shape [*";
        for (auto size : normalized_shape) {
            ss << ", " << size;
        }
        ss << "], but got input of size" << input_shape;
        throw std::runtime_error(ss.str());
    }

    compute_n1_n2(input, normalized_shape, n1, n2);
}

void check_args(at::Tensor input, at::IntArrayRef normalized_shape, at::Tensor gamma,
                at::Tensor beta, int& n1, int& n2) {
    check_args(input, normalized_shape, n1, n2);
    check_args(normalized_shape, gamma, beta);
}

void cuda_layer_norm(at::Tensor* output, at::Tensor* mean, at::Tensor* invvar, at::Tensor* input,
                     int n1, int n2, at::IntArrayRef normalized_shape, at::Tensor* gamma,
                     at::Tensor* beta, double epsilon);

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) \
    CHECK_CUDA(x);     \
    CHECK_CONTIGUOUS(x)

std::vector<at::Tensor> layer_norm_affine(at::Tensor input, at::IntArrayRef normalized_shape,
                                          at::Tensor gamma, at::Tensor beta, double epsilon) {
    CHECK_INPUT(input);
    CHECK_INPUT(gamma);
    CHECK_INPUT(beta);
    int n1, n2;
    check_args(input, normalized_shape, gamma, beta, n1, n2);

    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));

    at::Tensor output = at::empty_like(input, gamma.options().dtype(gamma.scalar_type()));
    at::Tensor mean = at::empty({n1}, input.options().dtype(at::ScalarType::Float));
    at::Tensor invvar = at::empty_like(mean);

    cuda_layer_norm(&output, &mean, &invvar, &input, n1, n2, normalized_shape, &gamma, &beta,
                    epsilon);

    return {output, mean, invvar};
}

void cuda_layer_norm_gradient(at::Tensor* dout, at::Tensor* mean, at::Tensor* invvar,
                              at::Tensor* input, int n1, int n2, at::IntArrayRef normalized_shape,
                              at::Tensor* gamma, at::Tensor* beta, double epsilon,
                              at::Tensor* grad_input, at::Tensor* grad_gamma,
                              at::Tensor* grad_beta);

std::vector<at::Tensor> layer_norm_gradient_affine(at::Tensor dout, at::Tensor mean,
                                                   at::Tensor invvar, at::Tensor input,
                                                   at::IntArrayRef normalized_shape,
                                                   at::Tensor gamma, at::Tensor beta,
                                                   double epsilon) {
    CHECK_INPUT(dout);
    CHECK_INPUT(mean);
    CHECK_INPUT(invvar);
    CHECK_INPUT(input);
    CHECK_INPUT(gamma);
    CHECK_INPUT(beta);
    int n1, n2;
    check_args(input, normalized_shape, gamma, beta, n1, n2);

    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));

    at::Tensor grad_input = at::empty_like(input);
    at::Tensor grad_gamma = at::empty_like(gamma);
    at::Tensor grad_beta = at::empty_like(beta);

    cuda_layer_norm_gradient(&dout, &mean, &invvar, &input, n1, n2, normalized_shape, &gamma, &beta,
                             epsilon, &grad_input, &grad_gamma, &grad_beta);

    return {grad_input, grad_gamma, grad_beta};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward_affine", &layer_norm_affine, "LayerNorm forward (CUDA)");
    m.def("backward_affine", &layer_norm_gradient_affine, "LayerNorm backward (CUDA)");
}